#include "dht.h"
#include "bencode.h"
#include "bt.h"
#include "epoll_loop.h"
#include "fstream.h"
#include "http.h"
#include "krpc.h"
#include "routing_table.h"
#include "socket.h"
#include "task.h"
#include <bitset>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <functional>
#include <iostream>
#include <mysql-cppconn/mysqlx/xdevapi.h>
#include <openssl/sha.h>
#include <optional>
#include <random>
#include <sstream>
#include <string>
#include <string_view>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std::chrono_literals;
namespace cobt {
DHTDigger::DHTDigger(mysqlx::Table &table) : m_table(table) {
  m_id.resize(SHA_DIGEST_LENGTH);
  SHA1((const unsigned char *)"COBT_YYYY_2024", 15,
       (unsigned char *)m_id.data());
  m_rt.m_k = 8;
  m_rt.m_node.m_id = DHTNode::ConvertToBitset(m_id);
  m_rt.m_node.m_id_str = m_id;
}
void DHTDigger::set_info_hash(const std::string &info_hash, bool status,
                              const std::string &name) {
  std::stringstream ss;
  for (auto ch : info_hash) {
    ss << std::hex << std::setw(2) << std::setfill('0')
       << (unsigned int)(unsigned char)ch;
  }
  mysqlx::RowResult row_result = m_table.select("info_hash")
                                     .where("info_hash = :hash")
                                     .bind("hash", ss.str())
                                     .execute();

  if (row_result.fetchOne()) {
    // 如果 info_hash 已存在，执行更新操作
    m_table.update()
        .set("status", status)
        .set("name", name)
        .where("info_hash = :hash")
        .bind("hash", ss.str())
        .execute();
  } else {
    // 如果 info_hash 不存在，执行插入操作
    m_table.insert("info_hash", "status", "name")
        .values(ss.str(), status, name)
        .execute();
  }
}

bool DHTDigger::info_hash_status(const std::string &info_hash) const {
  std::stringstream ss;
  for (auto ch : info_hash) {
    ss << std::hex << std::setw(2) << std::setfill('0')
       << (unsigned int)(unsigned char)ch;
  }
  auto row_result = m_table.select("status")
                        .where("info_hash = :info_hash")
                        .bind("info_hash", ss.str())
                        .execute();
  auto row = row_result.fetchOne();
  if (!row) {
    return false;
  }
  return row[0].get<bool>();
}

Task<> DHTDigger::get_metadata(DHTNode node) {
  auto res = co_await GetMetaFile(node.m_addr, node.m_info_hash, m_id)
                 .get_meta_file(*this);
  --COBT_TCP_CONNECTIONS;
  if (!res.empty() && !info_hash_status(node.m_info_hash)) {
    set_info_hash(node.m_info_hash, true, res);
  }
}

void DHTDigger::get_metadata_loop() {
  int q_size = m_peers_queue.size();
  std::cout << "q_size: " << q_size
            << " COBT_TCP_CONNECTIONS: " << COBT_TCP_CONNECTIONS << std::endl;
  std::unordered_map<std::string, bool> temp_status;
  for (int i = 0; i < q_size; i++) {
    auto info_hash = m_peers_queue.front().m_info_hash;
    bool status = false;
    if (temp_status.contains(info_hash)) {
      status = temp_status[info_hash];
    } else {
      status = info_hash_status(info_hash);
      temp_status[info_hash] = status;
    }
    if (status) {
      m_peers_queue.pop();
      continue;
    }
    if (COBT_TCP_CONNECTIONS > 512) {
      break;
    }
    ++COBT_TCP_CONNECTIONS;
    loop_enqueue(&DHTDigger::get_metadata, this,
                 std::move(m_peers_queue.front()));
    m_peers_queue.pop();
    continue;
  }
  set_timeout(1s, [this]() { get_metadata_loop(); });
}

std::string DHTDigger::insert_id(const SocketAddress &addr,
                                 const BEncodeData &a) {
  auto id = a["id"];
  if (!id) {
    return "";
  }
  auto &id_str = *id->as<std::string>();
  if (id_str.size() != ID_LENGTH / 8) {
    return "";
  }
  if (addr == m_rt.m_node.m_addr) {
    return "";
  }
  // DHTNode
  m_rt.insert(DHTNode(id_str, addr));
  return id_str;
}

std::optional<BMap> DHTDigger::get_closet_nodes(const std::string &id) {
  if (id.size() != ID_LENGTH / 8) {
    return std::nullopt;
  }
  auto nodes = m_rt.closest_nodes(id);
  if (nodes.size() < m_rt.m_k) {
    nodes.push_back(m_rt.m_node);
  } else {
    nodes.back() = m_rt.m_node;
  }
  std::string nodes4;
  std::string nodes6;
  for (const auto &node : nodes) {
    if (node.m_addr.get_family() == AF_INET) {
      nodes4.append(node.compact_info());
    } else {
      nodes6.append(node.compact_info());
    }
  }

  BMap response{{"id", BEncodeData(m_id)}};
  if (!nodes4.empty()) {
    response["nodes"] = BEncodeData(std::move(nodes4));
  }
  if (!nodes6.empty()) {
    response["nodes6"] = BEncodeData(std::move(nodes6));
  }
  return response;
}

void DHTDigger::insert_peer(DHTNode node) {
  if (m_peers_queue.size() > MAX_QUEUE_SIZE) {
    return;
  }
  m_peers_queue.emplace(std::move(node));
}

void DHTDigger::handle_ping(const BEncodeData &t, const BEncodeData &a,
                            const SocketAddress &addr) {
  auto id = insert_id(addr, a);
  if (id.empty()) {
    return;
  }
  id = GenFakeId(id);
  m_send_queue.push([id, fd = m_fd, addr = addr, t = t]() {
    BEncodeData resp(BMap{{"t", t},
                          {"y", BEncodeData("r")},
                          {"r", BMap{{"id", BEncodeData(id)}}}});
    std::string send_buf = resp.dump();
    sendto(fd, send_buf, addr, 0);
  });
}
void DHTDigger::handle_find_node(const BEncodeData &t, const BEncodeData &a,
                                 const SocketAddress &addr) {
  auto id = insert_id(addr, a);
  if (id.empty()) {
    return;
  }
  m_rt.m_node.m_id_str = GenFakeId(id);
  m_id = m_rt.m_node.m_id_str;
  auto target = a["target"];
  if (!target) {
    return;
  }
  auto resp_r = get_closet_nodes(*target->as<std::string>());
  if (!resp_r) {
    return;
  };
  m_send_queue.push([id, fd = m_fd, addr = addr,
                     resp_r = BEncodeData(std::move(resp_r.value())), t = t]() {
    BEncodeData resp(BMap{{"t", t}, {"y", BEncodeData("r")}, {"r", resp_r}});
    std::string send_buf = resp.dump();
    sendto(fd, send_buf, addr, 0);
  });
}

void DHTDigger::handle_get_peers(const BEncodeData &t, const BEncodeData &a,
                                 const SocketAddress &addr) {
  auto id = insert_id(addr, a);
  if (id.empty()) {
    return;
  }
  m_rt.m_node.m_id_str = GenFakeId(id);
  m_id = m_rt.m_node.m_id_str;
  auto info_hash = a["info_hash"];
  if (!info_hash) {
    return;
  }
  const std::string &info_hash_str = *info_hash->as<std::string>();
  if (info_hash_str.size() != ID_LENGTH / 8) {
    return;
  }
  auto resp_r = get_closet_nodes(info_hash_str);
  if (!resp_r) {
    return;
  };

  std::string token = GenToken(info_hash_str);
  resp_r.value()["token"] = BEncodeData(token);
  m_send_queue.push([id, fd = m_fd, addr = addr,
                     resp_r = BEncodeData(std::move(resp_r.value())), t = t]() {
    BEncodeData resp(BMap{{"t", t}, {"y", BEncodeData("r")}, {"r", resp_r}});
    std::string send_buf = resp.dump();
    sendto(fd, send_buf, addr, 0);
  });
  if (info_hash_status(info_hash_str)) {
    return;
  }
  BEncodeData req{BMap{{"t", BEncodeData("aa")}, {"y", BEncodeData("q")}, {}}};
  auto nodes = m_rt.closest_nodes(info_hash_str);
  for (auto &node : nodes) {
    m_info_hashs[sockaddr_to_string(node.m_addr)] = info_hash_str;
    m_send_queue.push([id, fd = m_fd, addr = std::move(node.m_addr),
                       info_hash = info_hash_str]() {
      ::cobt::get_peers(fd, addr, GenFakeId(id), info_hash);
    });
  }
}

void DHTDigger::handle_announce_peer(const BEncodeData &t, const BEncodeData &a,
                                     const SocketAddress &addr) {
  auto id = insert_id(addr, a);
  if (id.empty()) {
    return;
  }
  m_rt.m_node.m_id_str = GenFakeId(id);
  m_id = m_rt.m_node.m_id_str;
  auto info_hash = a["info_hash"];
  if (!info_hash) {
    return;
  }
  auto &info_hash_str = *info_hash->as<std::string>();
  if (info_hash_str.size() != ID_LENGTH / 8) {
    return;
  }

  m_send_queue.push([t, fd = m_fd, addr = addr, id = m_id]() {
    BEncodeData resp(BMap{{"t", t},
                          {"y", BEncodeData("r")},
                          {"r", BMap{{"id", BEncodeData(id)}}}});
    std::string send_buf = resp.dump();
    sendto(fd, send_buf, addr, 0);
  });

  if (info_hash_status(info_hash_str)) {
    return;
  }
  if (m_peers_queue.size() >= MAX_QUEUE_SIZE) {
    return;
  }

  if (auto iport = a["implied_port"]; iport) {
    auto i = iport->as<int64_t>();
    if (!i || *i != 0) {
      return;
    }
  }
  auto pport = a["port"];
  if (!pport) {
    return;
  }
  if (auto port = pport->as<int64_t>(); port) {
    SocketAddress naddr = addr;
    naddr.set_port(*port);
    // std::cout << "handle_announce_peer from " << addr << std::endl;
    DHTNode node(id, std::move(naddr), true);
    node.m_info_hash = info_hash_str;
    insert_peer(std::move(node));
  }
  // set_info_hash(info_hash_str, false);
}

void DHTDigger::handle_find_node_response(const std::string &id,
                                          const BEncodeData &r,
                                          const std::string &key, int step) {
  auto resp_nodes = r[key];
  if (!resp_nodes) {
    return;
  }
  auto resp_nodes_str = resp_nodes->as<std::string>();
  if (!resp_nodes_str || resp_nodes_str->size() % step) {
    return;
  }
  for (int i = 0; i < resp_nodes_str->size(); i += step) {
    auto node = parse_compact(resp_nodes_str->substr(i, step)).value();
    m_send_queue.push([id, fd = m_fd, addr = std::move(node.m_addr)]() {
      ping(fd, addr, GenFakeId(id));
    });
    break;
  }
}

void DHTDigger::handle_get_peers_response_node(const std::string &id,
                                               const BEncodeData &r,
                                               const std::string &key, int step,
                                               const SocketAddress &addr) {
  auto resp_nodes = r[key];
  if (!resp_nodes) {
    return;
  }
  auto resp_nodes_str = resp_nodes->as<std::string>();
  if (!resp_nodes_str || resp_nodes_str->size() % step) {
    return;
  }
  auto addr_key = sockaddr_to_string(addr);
  if (!m_info_hashs.contains(addr_key)) {
    return;
  }
  auto info_hash = m_info_hashs[sockaddr_to_string(addr)];
  if (info_hash_status(info_hash)) {
    return;
  }
  m_info_hashs.erase(addr_key);
  for (int i = 0; i < resp_nodes_str->size(); i += step) {
    auto node = parse_compact(resp_nodes_str->substr(i, step)).value();
    auto key = sockaddr_to_string(node.m_addr);
    if (m_info_hashs.contains(key)) {
      continue;
    }
    if (m_info_hashs.size() > MAX_INFO_HASH_SIZE) {
      return;
    }
    m_info_hashs[key] = info_hash;
    m_send_queue.push(
        [id, fd = m_fd, addr = std::move(node.m_addr), info_hash]() {
          ::cobt::get_peers(fd, addr, GenFakeId(id), info_hash);
        });

    // break;
  }
}

void DHTDigger::handle_get_peers_response_peer(const std::string &id,
                                               const BEncodeData &values,
                                               const SocketAddress &addr) {
  auto resp_val_list = values.as<std::vector<BEncodeData>>();
  if (!resp_val_list) {
    return;
  }
  auto addr_key = sockaddr_to_string(addr);
  if (!m_info_hashs.contains(addr_key)) {
    std::cout << "handle_get_peers_response_peer not exist" << std::endl;
    return;
  }
  auto info_hash = m_info_hashs[sockaddr_to_string(addr)];
  if (info_hash_status(info_hash)) {
    std::cout << "handle_get_peers_response_peer exist" << std::endl;
    return;
  }
  std::cout << "handle_get_peers_response_peer" << std::endl;

  m_info_hashs.erase(addr_key);
  for (const auto &peer : *resp_val_list) {
    auto peer_str = peer.as<std::string>();
    if (!peer_str) {
      continue;
    }
    auto peer_addr = parse_compact_addr(*peer_str);
    if (!peer_addr) {
      continue;
    }
    DHTNode peer_node(std::bitset<ID_LENGTH>(0), *peer_addr);
    peer_node.m_peer = true;
    insert_peer(std::move(peer_node));
  }
}

void DHTDigger::handle_response(const BEncodeData &bdata,
                                const SocketAddress &addr) {
  auto r = bdata["r"];
  if (!r) {
    return;
  }
  std::string id = insert_id(addr, *r);
  if (id.empty()) {
    return;
  }
  auto token = (*r)["token"];
  if (!token) {
    if (!(*r)["nodes"] && !(*r)["nodes6"]) {
      m_send_queue.push([id, fd = m_fd, addr = addr]() {
        find_node(fd, addr, GenFakeId(id), GenToken(id, id.size()));
      });
    }
    handle_find_node_response(id, *r, "nodes", 26);
    handle_find_node_response(id, *r, "nodes6", 38);
  } else {
    handle_get_peers_response_node(id, *r, "nodes", 26, addr);
    handle_get_peers_response_node(id, *r, "nodes6", 38, addr);
  }
  auto values = (*r)["values"];
  if (values) {
    handle_get_peers_response_peer(id, *values, addr);
  }
}

void DHTDigger::handle_request(const BEncodeData &bdata,
                               const SocketAddress &addr) {
  auto t = bdata["t"];
  if (!t) {
    return;
  }
  auto y = bdata["y"];
  if (!y || *(y->as<std::string>()) != "q") {
    return;
  }
  auto a = bdata["a"];
  if (!a) {
    return;
  }
  auto q = bdata["q"];
  if (!q) {
    return;
  }
  auto q_str = *(q->as<std::string>());
  if (q_str == "ping") {
    handle_ping(*t, *a, addr);
    return;
  }
  if (q_str == "get_peers") {
    handle_get_peers(*t, *a, addr);
    return;
  }
  if (q_str == "find_node") {
    handle_find_node(*t, *a, addr);
    return;
  }
  if (q_str == "announce_peer") {
    handle_announce_peer(*t, *a, addr);
    return;
  }
  return;
}

void DHTDigger::handle_recv(std::span<char> buf, const SocketAddress &addr) {
  std::string buf_str(buf.begin(), buf.end());
  auto bdata = BEncodeData::parse(buf_str);
  if (!bdata) {
    return;
  }
  auto y = bdata.value()["y"];
  if (!y) {
    return;
  }
  auto y_str = y->as<std::string>();
  if (!y_str) {
    return;
  }

  if (*y_str == "q") {
    handle_request(bdata.value(), addr);
  } else {
    handle_response(bdata.value(), addr);
  }
}

std::string DHTDigger::GenToken(const std::string &info_hash, int len) {
  std::random_device rd;
  std::mt19937 generator(rd());
  std::uniform_int_distribution<> distribution(0, info_hash.size() - 1);

  std::string randomString;
  randomString.reserve(len);

  for (size_t i = 0; i < len; ++i) {
    randomString += info_hash[distribution(generator)];
  }
  return randomString;
}

std::string DHTDigger::GenFakeId(const std::string &id) {
  if (id.empty()) {
    return "";
  }
  std::string res = id;
  auto id_size = id.size();
  std::string fake_id = GenToken(id);
  for (int i = 0; i < TOKEN_SIZE; i++) {
    res[id_size - i - 1] = fake_id[i];
  }
  return res;
}

void DHTDigger::join() {
  set_timeout(1s, [this]() { join(); });
  auto nodes = m_rt.closest_nodes(GenToken(m_id, m_id.size()));
  if (nodes.size() < 2) {
    m_send_queue.push([fd = m_fd, id = m_id]() {
      ping(fd, SocketAddress{"dht.transmissionbt.com", 6881}, GenFakeId(id));
    });
  } else {
    m_send_queue.push([fd = m_fd, addr = nodes[0].m_addr, id = m_id]() {
      ping(fd, addr, GenFakeId(id));
    });
  }
}

Task<> DHTDigger::init() {
  HttpClient hc("test.ipw.cn", 80);

  auto resp = co_await hc.get("/");
  if (resp) {
    std::cout << resp->get_body() << std::endl;
    m_rt.m_node.m_addr = SocketAddress(resp->get_body(), 6881);
  }
  auto t =
      recvfrom_async(m_fd, 0,
                     std::bind(&DHTDigger::handle_recv, this,
                               std::placeholders::_1, std::placeholders::_2));
  g_loop.enqueue(t.m_handle);
  t.m_handle = nullptr;
  join();
  get_metadata_loop();
}

void DHTDigger::listen(uint16_t port) {
  int fd = socket(AF_INET6, SOCK_DGRAM, 0);
  if (fd < 0) {
    std::cout << "socket error" << std::endl;
    return;
  }
  m_fd = fd;
  int opt = 0;
  setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt));

  sockaddr_in6 addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin6_family = AF_INET6;
  addr.sin6_addr = in6addr_any;
  addr.sin6_port = htons(port);

  if (bind(fd, (sockaddr *)&addr, sizeof(addr)) < 0) {
    if (errno == EINPROGRESS) {
      std::cout << "wait bind in progress" << std::endl;
    }
    std::cout << "wait bind error:" << errno << std::endl;
    return;
  }

  int flags = fcntl(fd, F_GETFL, 0);
  fcntl(fd, F_SETFL, flags | O_NONBLOCK);

  loop_enqueue(&DHTDigger::init, this);

  while (!cobt::g_loop.empty()) {
    try {
      cobt::g_loop.run();
    } catch (std::exception &e) {
      std::cout << "exception:" << e.what() << std::endl;
      continue;
    }
  }
}

} // namespace cobt